use alloc::string::String;
use hashbrown::HashMap;

use crate::ir::{graph::Graph, tensor::Tensor};
use crate::ops::{
    context::OpsContext,
    register::{DeviceKind, get_ops},
};

pub fn run(graph: &Graph, mut feeds: HashMap<String, Tensor>) -> HashMap<String, Tensor> {
    let mut vals: HashMap<String, Tensor> = HashMap::new();
    for v in &graph.values {
        if let Some(c) = &v.constant {
            vals.insert(v.name.clone(), c.clone());
        }
    }

    for (k, v) in feeds.drain() {
        vals.insert(k, v);
    }

    for n in &graph.nodes {
        // skip no-op nodes
        if n.op_type.is_empty() {
            continue;
        }

        // skip if outputs already ready
        let outputs_already_ready = n.outputs.iter().all(|&vid| {
            let name = &graph.values[vid as usize].name;
            vals.contains_key(name)
        });
        if outputs_already_ready {
            continue;
        }

        let op_name = n.op_type.as_str();

        // gather inputs
        let mut inputs: alloc::vec::Vec<Tensor> = alloc::vec::Vec::with_capacity(n.inputs.len());
        for &vid in &n.inputs {
            let name = &graph.values[vid as usize].name;
            let t = vals
                .get(name)
                .unwrap_or_else(|| panic!("missing input tensor: {}", name));
            inputs.push(t.clone());
        }

        // choose execution dtype from first input if available, else fallback to first constant output type later
        let exec_dtype = if let Some(t0) = inputs.get(0) {
            t0.dtype
        } else {
            panic!("no inputs for op {}", op_name)
        };

        // prepare outputs (preallocated with first input shape to avoid reallocation on common path)
        let mut outputs: alloc::vec::Vec<Tensor> = alloc::vec::Vec::with_capacity(n.outputs.len());
        let default_shape = if let Some(t0) = inputs.get(0) {
            t0.shape.clone()
        } else {
            alloc::vec::Vec::new()
        };
        for _ in &n.outputs {
            outputs.push(Tensor::new(exec_dtype, default_shape.clone()));
        }

        // collect attributes (values only)
        let attrs: alloc::vec::Vec<crate::ir::attribute::AttrValue> =
            n.attrs.iter().map(|(_, v)| v.clone()).collect();

        // fetch op from registry and compute
        let op = get_ops(op_name, exec_dtype, &DeviceKind::cpu0(), (1, 0))
            .unwrap_or_else(|| panic!("unsupported op or dtype: {} ({:?})", op_name, exec_dtype));

        let mut context = OpsContext {
            inputs,
            outputs,
            attrs,
        };
        op.compute(&mut context).expect("op compute failed");

        // write back outputs
        for (oi, &vid) in n.outputs.iter().enumerate() {
            let name = graph.values[vid as usize].name.clone();
            vals.insert(name, context.outputs[oi].clone());
        }
    }

    let mut out_map: HashMap<String, Tensor> = HashMap::new();
    for &vid in &graph.outputs {
        let name = graph.values[vid as usize].name.clone();
        out_map.insert(name.clone(), vals.remove(&name).unwrap());
    }
    out_map
}
